# 数据集相关路由
import os
import io

from exts import db
from flask import request, jsonify, g, send_file
from flask_sse import sse
from models import User, Dataset, Plant, Model
import threading
from ACRmodel.preprocess import preprocess_bed_to_npy
import shutil
from pathlib import Path


def register_dataset_routes(app):
    # 查询所有数据集
    @app.route('/plantACR/getDatasets', methods=['GET'])
    def get_datasets():
        resultVO = {
            "code": 200,
            "message": "成功获取数据集列表",
            "data": {
                "datasets": []
            }
        }
        try:
            # 按ID升序查询所有数据集
            datasets = Dataset.query.order_by(Dataset.id.asc()).all()
            if datasets:
                for dataset in datasets:
                    dataset_info = {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                        "plant_name": dataset.plant_name,
                        "user_name": dataset.user_name,
                        "chr_count": dataset.chr_count,
                        "chr_all": dataset.chr_all,
                        "sample_count": dataset.sample_count,
                        "description": dataset.description,
                        "files": dataset.files,
                        "file_size": dataset.file_size,
                        "public": dataset.public,
                        "status": dataset.status,
                        "create_time": dataset.create_time,
                        "update_time": dataset.update_time
                    }
                    resultVO["data"]["datasets"].append(dataset_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "数据集列表为空"
                return jsonify(resultVO)
        except Exception as e:
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取数据集列表失败！" + str(e)
            return jsonify(resultVO)

    # 查询所有公开数据集
    @app.route('/plantACR/getPublicDatasets', methods=['GET'])
    def get_my_datasets():
        resultVO = {
            "code": 200,
            "message": "成功获取公开数据集列表",
            "data": {
                "datasets": []
            }
        }
        try:
            # 按ID升序查询所有公开数据集
            datasets = Dataset.query.filter_by(status='公开').order_by(Dataset.id.asc()).all()
            if datasets:
                for dataset in datasets:
                    dataset_info = {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                        "plant_name": dataset.plant_name,
                        "user_name": dataset.user_name,
                        "chr_count": dataset.chr_count,
                        "chr_all": dataset.chr_all,
                        "sample_count": dataset.sample_count,
                        "description": dataset.description,
                        "files": dataset.files,
                        "file_size": dataset.file_size,
                        "public": dataset.public,
                        "status": dataset.status,
                        "create_time": dataset.create_time,
                        "update_time": dataset.update_time
                    }
                    resultVO["data"]["datasets"].append(dataset_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "公开数据集列表为空"
                return jsonify(resultVO)
        except Exception as e:
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取公开数据集列表失败！" + str(e)
            return jsonify(resultVO)

    # 查询所有个人数据集
    @app.route('/plantACR/getMyDatasets', methods=['GET'])
    def get_public_datasets():
        resultVO = {
            "code": 200,
            "message": "成功获取个人开数据集列表",
            "data": {
                "datasets": []
            }
        }
        try:
            # 按ID升序查询所有公开数据集
            datasets = Dataset.query.filter_by(user_id=g.id).order_by(Dataset.id.asc()).all()
            if datasets:
                for dataset in datasets:
                    dataset_info = {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                        "plant_name": dataset.plant_name,
                        "user_name": dataset.user_name,
                        "chr_count": dataset.chr_count,
                        "chr_all": dataset.chr_all,
                        "sample_count": dataset.sample_count,
                        "description": dataset.description,
                        "files": dataset.files,
                        "file_size": dataset.file_size,
                        "public": dataset.public,
                        "status": dataset.status,
                        "create_time": dataset.create_time,
                        "update_time": dataset.update_time
                    }
                    resultVO["data"]["datasets"].append(dataset_info)
                return jsonify(resultVO)
            else:
                resultVO["message"] = "个人数据集列表为空"
                return jsonify(resultVO)
        except Exception as e:
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，获取个人数据集列表失败！" + str(e)
            return jsonify(resultVO)

    # 更新数据集信息
    @app.route('/plantACR/updateDataset', methods=['POST'])
    def update_dataset():
        resultVO = {
            "code": 200,
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            dataset_id = data.get('id')
            new_dataset_name = data.get('dataset_name')

            # 检查新数据集名字是否已存在（自己除外）
            existing_dataset = Dataset.query.filter_by(dataset_name=new_dataset_name).first()
            if existing_dataset and existing_dataset.id != dataset_id:
                resultVO["code"] = 400
                resultVO["message"] = "设置的数据集名称在数据库中已存在"
                return jsonify(resultVO)

            dataset = Dataset.query.get(dataset_id)

            if dataset:
                dataset.dataset_name = data.get('dataset_name', dataset.dataset_name)
                dataset.description = data.get('description', dataset.description)
                db.session.commit()
                resultVO["code"] = 200
                resultVO["data"] = {
                    "dataset": {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                        "description": dataset.description,
                    }
                }
                resultVO["message"] = "数据集信息更新成功"
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器错误，数据集更新失败！" + str(e)

    # 更新数据集公开状态
    @app.route('/plantACR/updateDatasetPublic', methods=['POST'])
    def update_dataset_status():
        resultVO = {
            "code": 200,
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            dataset_id = data.get('id')
            dataset = Dataset.query.get(dataset_id)
            if dataset:
                # 状态校验逻辑
                if dataset.status in ('未处理', '处理中', '处理失败'):
                    resultVO["code"] = 400
                    resultVO["message"] = "数据集未处理，无法更新公开状态"
                    return jsonify(resultVO)

                # 公开状态取反
                dataset.public = not dataset.public
                dataset.status = '公开' if dataset.public else '私密'  # 状态同步
                db.session.commit()
                print("dataset.public:", dataset.public)
                print("dataset.status:", dataset.status)

                resultVO["code"] = 200
                resultVO["data"] = {
                    "dataset": {
                        "id": dataset.id,
                        "public": dataset.public  # 返回字段为public
                    }
                }
                resultVO["message"] = "数据集状态更新成功"
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器错误，数据集公开状态更新失败！" + str(e)
            return jsonify(resultVO)

    # 删除数据集
    @app.route('/plantACR/deleteDataset', methods=['POST'])
    def delete_dataset():
        resultVO = {
            "code": 200,
            "data": {}
        }
        try:
            current_plant_role = g.role
            data = request.get_json()  # 从前端获取 JSON 数据
            dataset_id = data.get('id')
            dataset = Dataset.query.get(dataset_id)

            # 检查权限，只有管理员或用户自己能删除数据集
            if current_plant_role == 'admin' and g.id != dataset.user_id:
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            # 检查数据集是否存在关联模型
            dataset_models = Model.query.filter_by(dataset_id=dataset_id).first()
            if dataset_models:
                resultVO["code"] = 400
                resultVO["message"] = '该数据集存在关联模型，无法删除'
                return jsonify(resultVO)

            if dataset:
                # 删除本地数据集文件
                file_path = f"ACRmodel/static/dataset/{dataset.plant_name}/{dataset.id}"
                if os.path.exists(file_path):
                    shutil.rmtree(file_path)
                # 删除temp文件夹中的文件
                if dataset.files:
                    temp_dir = Path(app.root_path) / "ACRmodel" / "static" / "temp"
                    for filename in dataset.files.split(','):
                        file_path = temp_dir / filename.strip()
                        if file_path.exists():
                            os.remove(file_path)
                db.session.delete(dataset)
                db.session.commit()

                resultVO["data"] = {
                    "dataset": {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                    }
                }
                resultVO["message"] = "数据集删除成功"
                return jsonify(resultVO)
            else:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)
        except Exception as e:
            print(str(e))
            db.session.rollback()
            resultVO["code"] = 500
            resultVO["message"] = "服务器错误，删除数据集失败！"
            return jsonify(resultVO)

    # 创建数据集任务(未处理)
    @app.route('/plantACR/createDataset', methods=['POST'])
    def create_dataset():
        resultVO = {
            "code": 200,
            "message": "已创建数据集任务",
            "data": {}
        }
        try:
            data = request.get_json()  # 从前端获取 JSON 数据
            user_id = g.id

            user = User.query.get(user_id)
            plant = Plant.query.filter_by(name=data.get('plant_name')).first()
            existing_dataset = Dataset.query.filter_by(dataset_name=data.get('dataset_name')).first()

            # 检查用户 是否存在
            if not user:
                resultVO["code"] = 404
                resultVO["message"] = "登录信息异常，用户不存在"
                return jsonify(resultVO)

            # 检查植物 是否存在
            if not plant:
                resultVO["code"] = 404
                resultVO["message"] = "所选植物在数据库中不存在"
                return jsonify(resultVO)

            # 检查数据集名字是否已存在
            if existing_dataset:
                resultVO["code"] = 400
                resultVO["message"] = "设置的数据集名称已存在，请更换名称"
                return jsonify(resultVO)

            # 创建新的数据集记录
            new_dataset = Dataset(
                user_id=user.id,
                user_name=user.name,
                plant_id=plant.id,
                plant_name=plant.name,
                dataset_name=data.get('dataset_name'),
                public=data.get('public'),
                description=data.get('description'),
                files=data.get('files'),
                status="未处理",
            )
            db.session.add(new_dataset)
            db.session.commit()
            resultVO["data"] = {
                "dataset": {
                    "id": new_dataset.id,
                    "dataset_name": new_dataset.dataset_name,
                    "user_name": new_dataset.user_name,
                    "plant_name": new_dataset.plant_name,
                    "status": new_dataset.status,
                    "create_time": new_dataset.create_time
                }
            }
            return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            print(str(e))
            resultVO["code"] = 500
            resultVO["message"] = "数据集任务添加失败"
            return jsonify(resultVO)

    # 预处理数据集
    @app.route('/plantACR/preprocessDataset', methods=['POST'])
    def preprocess_dataset():
        resultVO = {
            "code": 200,
            "message": "已创建数据集处理任务",
            "data": {}
        }
        data = request.get_json()  # 从前端获取 JSON 数据
        dataset_id = data.get('id')

        try:
            dataset = Dataset.query.get(dataset_id)
            # 检查数据集是否存在
            if not dataset:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)

            # 检查是否有其他数据集正在处理
            processing_dataset = Dataset.query.filter_by(status='处理中').first()
            if processing_dataset:
                resultVO["code"] = 400
                resultVO["message"] = f"数据集 {processing_dataset.dataset_name} 正在处理中，请等待完成"
                return jsonify(resultVO)

            # 检查权限，只有管理员或用户自己能处理数据集
            if g.id != dataset.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            if dataset.status in ('未处理', '处理失败'):
                dataset.status = '处理中'
                db.session.commit()
                resultVO["code"] = 200
                resultVO["message"] = "数据集开始处理"
                resultVO["data"] = {
                    "dataset": {
                        "id": dataset.id,
                        "dataset_name": dataset.dataset_name,
                    }
                }
                # 启动后台线程执行处理
                threading.Thread(target=start_processing, args=(dataset.id,), daemon=True).start()
                return jsonify(resultVO)
            # 未处理 和 处理失败 之外的状态是无法发送处理请求的，返回数据异常
            else:
                resultVO["code"] = 400
                resultVO["message"] = "数据异常！"
                return jsonify(resultVO)
        except Exception as e:
            db.session.rollback()
            dataset = Dataset.query.get(dataset_id)
            if dataset:
                dataset.status = '处理失败'
                db.session.commit()
            print("数据集处理任务启动失败:", e)
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，数据集处理失败！"
            return jsonify(resultVO)

    # 在后台线程中执行数据处理
    def start_processing(dataset_id):
        """"
            status：未处理，处理中，处理失败，公开，私密

            :param
                dataset_id: 数据集ID
            :return
                1. 处理成功：更新数据集状态为公开/私密，发送SSE消息
                2. 处理失败：更新数据集状态为处理失败，发送SSE消息
                3. 数据集不存在：发送SSE消息
                4. 服务器错误：更新数据集状态为处理失败，发送SSE消息
        """
        with app.app_context():
            try:
                dataset = Dataset.query.get(dataset_id)
                if dataset:
                    print("启动数据集处理...")
                    # 调用数据处理函数
                    result, mes, chrom_count, chrom_str, file_size_mb, sample_count = preprocess_bed_to_npy(
                        dataset.plant_name, dataset.id, dataset.files, 0)
                    if result:
                        print("处理成功！")
                        status = '公开' if dataset.public else '私密'  # 公开状态同步
                        dataset = Dataset.query.get(dataset_id)
                        dataset.chr_count = chrom_count
                        dataset.chr_all = chrom_str
                        dataset.status = status
                        dataset.file_size = file_size_mb
                        dataset.sample_count = sample_count
                        db.session.commit()
                        sse_dataset_status(dataset_id, status, mes)
                        if dataset.files:
                            # app.root_path = F:\Project\PlantACR
                            temp_dir = Path(app.root_path) / "ACRmodel" / "static" / "temp"
                            for filename in dataset.files.split(','):
                                file_path = temp_dir / filename.strip()
                                if file_path.exists():
                                    os.remove(file_path)
                                    print(f"已删除未处理文件: {file_path}")
                    else:
                        print("处理失败！")
                        sse_dataset_status(dataset_id, "处理失败", mes)
                else:
                    print("数据集处理任务不存在")
                    sse_dataset_status(dataset_id, "处理失败", "数据集处理任务不存在")
            except Exception as e:
                db.session.rollback()
                print("数据集处理失败：", e)
                sse_dataset_status(dataset_id, "处理失败", "服务器错误，处理失败！")
            finally:
                print("finally:当前数据集处理结束")

    # 发送SSE消息，告知前端数据集处理状态变化，并更新数据集在数据库中的状态
    def sse_dataset_status(dataset_id, status, sse_message=None):
        dataset = Dataset.query.get(dataset_id)
        if dataset:
            dataset.status = status
            db.session.commit()
            # 向所有订阅了 dataset_status 事件的客户端发送消息
            try:
                # 发送SSE消息
                sse.publish({
                    "dataset_id": dataset_id,
                    "dataset_name": dataset.dataset_name,
                    "status": status,
                    "message": sse_message
                }, type='dataset_status', channel='sse')
                print("SSE消息已发布")
            except Exception as e:
                print("SSE发布失败：", e)
            return "消息已发送"
        else:
            print("数据集不存在")
            # 发送SSE消息
            sse.publish({
                "dataset_id": dataset_id,
                "dataset_name": '',
                "status": status,
                "message": sse_message
            }, type='dataset_status', channel='sse')
            return "消息已发送，数据集不存在"

    # 下载数据集
    @app.route('/plantACR/downloadDataset/<int:dataset_id>', methods=['GET'])
    def download_datatest(dataset_id):
        resultVO = {
            "code": 200,
            "message": ""
        }
        try:
            dataset = Dataset.query.get(dataset_id)
            if not dataset:
                resultVO["code"] = 404
                resultVO["message"] = "数据集不存在"
                return jsonify(resultVO)  # 虽然是json，但是前端以二进制文件Blob接受，需要进一步读取

            # 公开数据集都可以下载，私密数据集只有管理员和用户自己可以下载
            if (not dataset.public) and g.id != dataset.user_id and g.role != 'admin':
                resultVO["code"] = 403
                resultVO["message"] = '权限不足！'
                return jsonify(resultVO)

            file_name = f"dataset_{dataset_id}.zip"
            dataset_file = Path(app.root_path) / "ACRmodel" / "static" / "dataset" / dataset.plant_name / str(
                dataset_id) / file_name

            # 文件不存在
            if not dataset_file.exists():
                resultVO["code"] = 404
                resultVO["message"] = "文件不存在"
                return jsonify(resultVO)

            # 获取二进制流
            with open(dataset_file, 'rb') as file:
                bytes_io = io.BytesIO(file.read())
            # 返回文件给前端
            return send_file(bytes_io,
                             mimetype='application/zip, application/octet-stream',
                             as_attachment=True,  # 强制浏览器以附件形式下载文件，而非直接打开
                             download_name=file_name  # 设置默认下载文件名
                             )
        except Exception as e:
            print(str(e))
            resultVO["code"] = 500
            resultVO["message"] = "服务器异常，下载失败！"
            return jsonify(resultVO)

    return app
